#include "UdpClient.h"
#include <clusnet/base/WeakCallback.h>

using namespace clusnet::net;

UdpClient::UdpClient(boost::asio::io_context& io, char const* host, char const* port, boost::asio::ip::udp::resolver::protocol_type const& proto)
    : io_(io),
      strand_(io),
      socket_(io, proto),
      session_(nullptr),
      close_(false),
      waitRecv_(false),
      writeCallback_(nullptr),
      messageCallback_(nullptr)
{
    boost::asio::ip::udp::resolver resolver(io_);
    remoteEndpoint_ = *(resolver.resolve(proto, host, port));
}

UdpClient::~UdpClient()
{
    socket_.close();
    session_.reset();
}

void UdpClient::safeAsyncWrite(UdpDataPacketPtr const& data)
{
    auto self(shared_from_this());
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this, self, data]
    {
        if (close_)
        {
            return;
        }
        if (session_ == nullptr)
        {
            session_ = std::make_shared<UdpSession>(io_, strand_, socket_, remoteEndpoint_, clusnet::base::makeWeakCallback(self, &UdpClient::handleWriteComplete));
        }
        session_->safeAsyncWrite(data);
    }));
}

void UdpClient::safeClose()
{
    auto self(shared_from_this());
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this, self]
    {
        if (close_)
        {
            return;
        }
        close_ = true;
        waitRecv_ = false;
        socket_.close();
        session_.reset();
    }));
}

void UdpClient::safeStartReceive()
{
    auto self(shared_from_this());
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this, self]
    {
        if (close_)
        {
            return;
        }
        if (session_ == nullptr)
        {
            waitRecv_ = true;
            safeAsyncWrite(std::make_shared<UdpDataPacket>(0));
            return;
        }
        socket_.async_receive_from(boost::asio::buffer(recvBuffer_), recvEndpoint_, boost::asio::bind_executor(strand_, [this, self](boost::system::error_code const& ec, std::size_t len)
        {
            if (ec)
            {
                LOG(WARNING) << "recv failed. " << ec.message();
                return;
            }
            if (recvEndpoint_ != remoteEndpoint_)
            {
                LOG(WARNING) << "recv message from " << recvEndpoint_.address() << ":" << recvEndpoint_.port();
                return;
            }
            assert(messageCallback_ != nullptr);
            UdpUniqDataPacketPtr msg(new UdpDataPacket(len));
            msg->append(recvBuffer_.data(), len);
            messageCallback_(session_, std::move(msg));
            safeStartReceive();
        }));
    }));
}

void UdpClient::handleWriteComplete(UdpSessionPtr& session, UdpDataPacketPtr const& data)
{
    if (writeCallback_)
    {
        writeCallback_(session, data);
    }
    if (!waitRecv_)
    {
        return;
    }
    auto self(shared_from_this());
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this, self]
    {
        if (waitRecv_)
        {
            waitRecv_ = false;
            safeStartReceive();
        }
    }));
}
